import numpy as np
import torch

from decision_transformer.training.trainer import Trainer


class SequenceTrainer(Trainer):

    def train_step(self):
        states, actions, rewards, dones, rtg, timesteps, attention_mask = self.get_batch(self.batch_size)

        action_target = torch.clone(actions)
        reward_target = torch.clone(rewards)
        state_target = torch.clone(states)

        state_preds, action_preds, reward_preds = self.model.forward(
            states, actions, rewards, rtg[:,:-1], timesteps, attention_mask=attention_mask,
        )

        act_dim = action_preds.shape[2]
        action_preds = action_preds.reshape(-1, act_dim)[attention_mask.reshape(-1) > 0]
        action_target = action_target.reshape(-1, act_dim)[attention_mask.reshape(-1) > 0]

        if 1:
            state_dim = state_target.shape[-1]
            state_preds = state_preds.reshape(-1, state_dim)[attention_mask.reshape(-1) > 0]
            state_target = state_target.reshape(-1, state_dim)[attention_mask.reshape(-1) > 0]

            reward_target = reward_target.reshape(-1, 1)[attention_mask.reshape(-1) > 0]
            reward_preds = reward_preds.reshape(-1, 1)[attention_mask.reshape(-1) > 0]

            # action_preds.shape
            # action_target.shape
            # state_preds.shape
            # reward_preds.shape
            # state_dim = state_preds.shape[2]
            # state_preds.reshape(-1, state_dim)[attention_mask.reshape(-1) > 0].shape
            # action_preds = state_preds.reshape(-1, state_dim)[attention_mask.reshape(-1) > 0]
            # action_target = action_target.reshape(-1, act_dim)[attention_mask.reshape(-1) > 0]

            loss = self.loss_fn(
                state_preds, action_preds, reward_preds,
                state_target, action_target, reward_target,
            )
        else:
            loss = self.loss_fn(
                None, action_preds, None,
                None, action_target, None,
            )

        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), .25)
        self.optimizer.step()

        with torch.no_grad():
            self.diagnostics['training/action_error'] = torch.mean((action_preds-action_target)**2).detach().cpu().item()

        return loss.detach().cpu().item()
